import os
from datetime import date
import glob
import logging
import pandas as pd
import pdfplumber
import matplotlib.pyplot as plt
import plotly.express as px
from plotly.graph_objs import Layout, Font
from wordcloud import WordCloud
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import re
import jieba


class MedicalFoodHandler:
    def __init__(self, food_path):
        self.logger = logging.getLogger(__name__)
        if not os.path.exists(food_path):
            raise FileNotFoundError(
                f"Path {food_path} does not exist. pls make sure the path is correct."
            )
        self.logger.info("setup food_path=%s", food_path)
        self.food_path = food_path
        self.data_path = os.path.join(self.food_path, "data")
        self.pdf_path = os.path.join(self.data_path, "Food_Manual")
        self.logger.info("pdf path is %s", self.pdf_path)

        self.nutritional_info_file_name = "result1.xlsx"
        self.data_file_name = "data.xlsx"
        current_date = date.today().strftime("%Y-%m-%d")
        self.result_path = os.path.join(self.food_path, "result", current_date)
        self.clear_result()
        self.pdf_files = []
        self.read_pdf_files()
        self.nutritional_info_table = None
        self.product_category_table = None
        plt.rcParams["font.family"] = "Simsong"
        self.layout = dict(
            font=dict(
                family="Simsong",  # Set the font family to Simsong
                size=14,  # Set the font size
                color="black",  # Set the font color
            ),
            title=dict(
                x=0.5,  # Center the title horizontally
                xanchor="center",  # Anchor the title at the center
                y=0.95,  # Position the title near the top
                yanchor="top",  # Anchor the title at the top
            ),
        )

    def preprocess_text(self, text):
        text = text.lower().strip()
        return text
    
    def chinese_tokenizer(self, text):
        return jieba.cut(text)

    def create_description(self, row):
        text = f"{row['产品类别']} {row['适用人群']}"
        return self.preprocess_text(text)

    def content_based_recommendation(
        self, customer_description, similarity_threshold=0.1, top_n=5
    ):
        products_df = self.product_category_table.copy(deep=True)

        descriptions = products_df.apply(self.create_description, axis=1)
        vectorizer = TfidfVectorizer(tokenizer=self.chinese_tokenizer)
        
        product_vectors = vectorizer.fit_transform(descriptions)
        
        customer_description = self.preprocess_text(customer_description)
        customer_vector = vectorizer.transform([customer_description])

        similarities = cosine_similarity(customer_vector, product_vectors)
        top_indices = [
            i for i, sim in enumerate(similarities[0]) if sim >= similarity_threshold
        ]

        # Filter indices based on similarity threshold

        top_indices = sorted(
            top_indices, key=lambda i: similarities[0][i], reverse=True
        )[:top_n]

        # Select top products
        top_products = products_df.iloc[top_indices]
        if len(top_products) <= 0:
            self.logger.info(f'针对"{customer_description}"的需求，不推荐使用任何产品')
            return

        self.logger.info(f'针对"{customer_description}"推荐以下产品:')
        for _, product in top_products.iterrows():
            self.logger.info("%s - %s", product["注册证号"], product["产品名称"])

    def plot_product_category_counts(self):
        count_name = "获批量"
        # Assuming you have a column named '产品类别' in your DataFrame
        if "产品类别" not in self.product_category_table.columns:
            raise ValueError("产品类别 column not found in DataFrame")

        # Group by '产品类别' and count the number of entries, then sort in descending order
        product_counts = (
            self.product_category_table.groupby("产品类别")
            .size()
            .reset_index(name=count_name)
            .sort_values(by=count_name, ascending=False)
        )

        plt.figure(figsize=(12, 8))
        bars = plt.bar(
            product_counts["产品类别"], product_counts[count_name], color="skyblue"
        )
        plt.xlabel("产品类别")
        plt.ylabel(count_name)
        plt.title("产品类别的特医食品获批量")
        plt.xticks(
            rotation=45, ha="right"
        )  # Rotate x-axis labels for better readability
        plt.tight_layout()  # Adjust layout to prevent clipping of labels
        for bar in bars:
            yval = bar.get_height()
            plt.text(
                bar.get_x() + bar.get_width() / 2,
                yval + 0.05,  # Adjust the vertical position of the text
                round(yval, 2),  # Display the value rounded to 2 decimal places
                ha="center",
                va="bottom",
            )
        plt.savefig(os.path.join(self.result_path, "2.3-result.png"))

    def plot_fat_protein_histograms(self):
        protein_name = "蛋白质(g)"
        fat_name = "脂肪(g)"
        # Assuming you have columns named '脂肪含量' and '蛋白质含量' in your DataFrame
        if protein_name not in self.nutritional_info_table.columns:
            raise ValueError(f"{protein_name} column not found in DataFrame")
        if fat_name not in self.nutritional_info_table.columns:
            raise ValueError(f"{fat_name} column not found in DataFrame")

        plt.figure(figsize=(12, 8))
        # Plot histogram for protein content
        n_protein, bins_protein, patches_protein = plt.hist(
            self.nutritional_info_table[protein_name],
            bins=30,
            alpha=0.5,
            color="green",
            label="蛋白质含量",
            edgecolor="black",
        )

        # Plot histogram for fat content
        n_fat, bins_fat, patches_fat = plt.hist(
            self.nutritional_info_table[fat_name],
            bins=30,
            alpha=0.5,
            color="blue",
            label="脂肪含量",
            edgecolor="black",
        )

        plt.xlabel("含量")
        plt.ylabel("频数")
        plt.title("特医食品脂肪和蛋白质含量的频数分布直方图")
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
        # Add text labels for protein content
        for i in range(len(n_protein)):
            plt.text(
                (bins_protein[i] + bins_protein[i + 1]) / 2,
                n_protein[i],
                str(int(n_protein[i])),
                ha="center",
                va="bottom",
            )

        # Add text labels for fat content
        for i in range(len(n_fat)):
            plt.text(
                (bins_fat[i] + bins_fat[i + 1]) / 2,
                n_fat[i],
                str(int(n_fat[i])),
                ha="center",
                va="bottom",
            )

        plt.savefig(os.path.join(self.result_path, "2.4-result.png"))

    def generate_wordcloud(self):
        target_column = "适用人群"
        if target_column not in self.product_category_table.columns:
            raise ValueError(f"{target_column} column not found in DataFrame")

        # Combine all the text from the '适用人群' column
        text = " ".join(self.product_category_table[target_column].astype(str))

        # Generate the word cloud
        wordcloud = WordCloud(
            font_path="Arial Unicode.ttf",
            background_color="white",
            max_words=10,
            contour_width=3,
            contour_color="red",
        ).generate(text)

        # Display the generated word cloud
        plt.figure(figsize=(10, 5))
        plt.imshow(wordcloud, interpolation="bilinear")
        plt.axis("off")  # Hide the axes
        plt.title("特医食品适用人群词云图")
        plt.savefig(os.path.join(self.result_path, "2.5-result.png"))

    def create_sunburst_chart(self):
        people_name = "适用人群类别"
        source_name = "产品来源"
        count_name = "数量"
        # Assuming you have a column named '适用人群类别' in your DataFrame
        if people_name not in self.product_category_table.columns:
            raise ValueError("适用人群类别 column not found in DataFrame")

        grouped = (
            self.product_category_table.groupby([people_name, source_name])
            .size()
            .reset_index(name=count_name)
        )
        fig = px.sunburst(
            grouped,
            path=[people_name, source_name],
            values=count_name,
            title="特医食品产品来源与适用人群类别分布",
        )
        fig.update_traces(
            textinfo="label+text+value",
            insidetextorientation="radial",
        )
        fig.update_layout(self.layout)
        fig.write_image(
            os.path.join(self.result_path, "2.2-result.png"), width=800, height=600
        )

    def group_by_year_and_from(self):
        year_name = "登记年份"
        source_name = "产品来源"
        count_name = "获批量"
        domestic_name = "国产"
        imported_name = "进口"
        grouped = (
            self.product_category_table.groupby([year_name, source_name])
            .size()
            .reset_index(name=count_name)
        )

        grouped = grouped.sort_values(by=year_name, ascending=True)
        domestic = grouped[grouped[source_name] == domestic_name]
        imported = grouped[grouped[source_name] == imported_name]

        plt.figure(figsize=(10, 6))

        plt.plot(
            domestic[year_name].astype(int),
            domestic[count_name].astype(int),
            label=domestic_name,
            marker="o",
        )
        for x, y in zip(
            domestic[year_name].astype(int), domestic[count_name].astype(int)
        ):
            plt.text(x, y, str(y), ha="center", va="bottom")

        plt.plot(
            imported[year_name].astype(int),
            imported[count_name].astype(int),
            label=imported_name,
            marker="o",
        )
        for x, y in zip(
            imported[year_name].astype(int), imported[count_name].astype(int)
        ):
            plt.text(x, y, str(y), ha="center", va="bottom")

        plt.xlabel(year_name)
        plt.ylabel(count_name)
        plt.title("特医食品获批量按登记年份和产品来源统计")
        # plt.xlim(int(min_year) - 1, None)
        plt.legend()
        plt.grid(True)
        plt.savefig(os.path.join(self.result_path, "2.1-result.png"))

    def preprocessing_product_category(self):
        original_df = pd.read_excel(os.path.join(self.data_path, self.data_file_name))
        # 用于存储提取的数据
        product_category_data = []
        tissue_status_data = []
        applicable_population_data = []
        warning_notes = []

        for pdf_file in self.pdf_files:
            file_path = os.path.join(self.pdf_path, pdf_file)
            with pdfplumber.open(file_path) as pdf:
                product_category = ""
                tissue_status = ""
                applicable_population = ""
                warning_note = ""
                for page in pdf.pages:
                    text = page.extract_text()
                    if text:
                        if "产品类别" in text:
                            start_index = text.find("产品类别")
                            end_index = text.find("\n", start_index)
                            product_category = text[
                                start_index + len("产品类别") + 1 : end_index
                            ].strip()
                        if "组织状态" in text:
                            start_index = text.find("组织状态")
                            end_index = text.find("\n", start_index)
                            tissue_status = (
                                text[start_index + len("组织状态") + 1 : end_index]
                                .strip()
                                .replace("。", "")
                            )
                        if "适用人群" in text:
                            start_index = text.find("适用人群")
                            end_index = text.find("\n", start_index)
                            applicable_population = text[
                                start_index + len("适用人群") + 1 : end_index
                            ].strip()
                        if "警示说明和注意事项" in text:
                            start_index = text.find("警示说明和注意事项")
                            warning_note = text[
                                start_index + len("警示说明和注意事项") + 1 :
                            ].strip()
                            # Remove extra whitespace and newlines
                            warning_note = re.sub(r"\s+", " ", warning_note)

                product_category_data.append(product_category)
                tissue_status_data.append(tissue_status)
                applicable_population_data.append(applicable_population)
                warning_notes.append(warning_note)

        # 处理数据列表，将空字符串替换为None（表示留空）
        product_category_data = [
            None if category == "" else category for category in product_category_data
        ]
        tissue_status_data = [
            None if status == "" else status for status in tissue_status_data
        ]
        applicable_population_data = [
            None if population == "" else population
            for population in applicable_population_data
        ]

        # 在原始DataFrame中新增三列
        original_df["产品类别"] = product_category_data
        original_df["组织状态"] = tissue_status_data
        original_df["适用人群"] = applicable_population_data

        # 保存更新后的DataFrame到result2.xlsx
        result_file_path = os.path.join(self.result_path, "1.2-result2.xlsx")
        original_df.to_excel(result_file_path, index=False)
        self.logger.info(
            f"Saved the food applicability information to {result_file_path}"
        )

        original_df["适用人群类别"] = original_df.apply(
            self.classify_applicable_population, axis=1
        )
        result_file_path_1_3 = os.path.join(self.result_path, "1.3-result2.xlsx")
        original_df.to_excel(result_file_path_1_3, index=False)

        original_df[["产品来源", "登记年份"]] = original_df.apply(
            self.classify_name_id, axis=1, result_type="expand"
        )
        result_file_path_1_4 = os.path.join(self.result_path, "1.4-result2.xlsx")
        original_df.to_excel(result_file_path_1_4, index=False)

        self.product_category_table = original_df
        self.product_category_table["注意事项"] = warning_notes

        final_result_file_path = os.path.join(self.result_path, "final_result.xlsx")
        self.product_category_table.to_excel(final_result_file_path, index=False)

    def classify_name_id(self, row):
        # 国食注字+TY+4位年号+4位顺序号，
        # 顺序号第 1 位数字为“5”表示该食品为进口产品，顺序号第 1 位数字为“0”表示该
        # 食品为国产产品；4 位年号为该食品的登记年份。
        registration_number = row["注册证号"]
        parts = registration_number.split("TY")
        if len(parts) != 2:
            raise ValueError("Invalid registration number format")
        year_type = parts[1]
        year = year_type[:4]
        sequence_number = year_type[4:]
        if sequence_number[0] == "5":
            product_type = "进口"
        elif sequence_number[0] == "0":
            product_type = "国产"
        else:
            raise ValueError("Invalid sequence number format")
        return product_type, year

    def classify_applicable_population(self, row):
        if (
            "0-12月龄" in row["适用人群"]
            or "早产/低出生体重婴儿" in row["适用人群"]
            or "婴儿" in row["适用人群"]
        ):
            return "特医婴配食品"
        else:
            return "1岁以上特医食品"

    def preprocessing_nutritional_from_pdfs(self):
        original_result_path = os.path.join(
            self.food_path, "data", self.nutritional_info_file_name
        )
        original_result = pd.read_excel(original_result_path)
        header_columns = [column for column in original_result.columns]
        self.logger.info(f"the nutritional infomation is: {header_columns}")

        for pdf_file in self.pdf_files:
            file_path = os.path.join(self.pdf_path, pdf_file)
            registration_number = os.path.splitext(pdf_file)[0]
            new_row = pd.Series(index=original_result.columns, dtype=object)
            new_row[original_result.columns[0]] = registration_number
            for column in original_result.columns[1:]:
                new_row[column] = 0

            with pdfplumber.open(file_path) as pdf:
                for page in pdf.pages:
                    tables = page.extract_tables()
                    for table in tables:
                        if len(table) < 2:
                            continue
                        df = pd.DataFrame(table[1:], columns=table[0])
                        if "营养成分" in df.columns and "每100kJ" in df.columns:
                            for nutrient in header_columns:
                                target_row = df[df["营养成分"] == nutrient]
                                if not target_row.empty:
                                    target_column = target_row["每100kJ"]
                                    if len(target_column) > 0:
                                        target_value = target_column.values[0]
                                        if target_value is None:
                                            target_value = 0
                                        float_value = 0
                                        try:
                                            float_value = float(target_value)
                                        except Exception as e:
                                            pass
                                        new_row[nutrient] = float_value
            original_result = pd.concat(
                [original_result, pd.DataFrame([new_row])], ignore_index=True
            )
        result_file_path = os.path.join(
            self.result_path, f"1.1-{self.nutritional_info_file_name}"
        )
        original_result = original_result.sort_values(by="蛋白质(g)", ascending=False)
        self.nutritional_info_table = original_result
        original_result.to_excel(result_file_path, index=False)
        self.logger.info(f"Saved the nutritional information to {result_file_path}")

    def read_pdf_files(self):
        self.pdf_files = [f for f in os.listdir(self.pdf_path) if f.endswith(".pdf")]
        self.logger.info(f"Got %d pdf files", len(self.pdf_files))

    def clear_result(self):
        if not os.path.exists(self.result_path):
            os.mkdir(self.result_path)
            return
        files = glob.glob(os.path.join(self.result_path, "*"))
        for file in files:
            try:
                os.remove(file)
                self.logger.info(f"Clearing file: {file}")
            except Exception as e:
                self.logger.error(f"Failed to remove file {file}: {e}")

    def get_nutritional_info(self):
        pass

    def __str__(self):
        return f"{self.name} - {self.price} - {self.quantity} - {self.expiration_date}"
